#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : project
@File : main_val.py
@Author : 薄荷糖
@Time : 2025/9/9 18:04  
@脚本说明 : 
"""

import argparse
from tools.mysql_val import run_mysql_for_host_or_cidr
from tools.redis_val import redis_run
from tools.ip_val import run_ip
from tools.ssh_val import run_ssh
from tools.subdomain_val import run_subdomain
from tools.port_val import scan_port_run

import logging
logging.getLogger("paramiko").setLevel(logging.WARNING)
# 创建命令行解析器对象
parser = argparse.ArgumentParser(description="智慧AI漏扫")
subparsers = parser.add_subparsers(help="选择要使用的功能", dest="sub_command")
# 创建MySQL的解析器对象
mysql_subparsers = subparsers.add_parser("mysql", help="MySQL相关漏扫")
mysql_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="MySQL主机地址")
mysql_subparsers.add_argument("--port", "-P", dest="port", type=int, default=3306, help="MySQL端口号")
mysql_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/user.dict", help="MySQL用户名")
mysql_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/passwd.dict", help="MySQL密码本")
# 创建Redis的解析器对象
redis_subparsers = subparsers.add_parser("redis", help="Redis相关漏扫")
redis_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="Redis主机地址")
redis_subparsers.add_argument("--port", "-P", dest="port", type=int, default=6379, help="Redis端口号")
redis_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/passwd.dict", help="Redis密码本")
# 创建IP的解析器对象
ip_subparsers = subparsers.add_parser("ip", help="主机相关漏扫")
ip_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="IP主机地址网段")
ip_subparsers.add_argument("--iface", "-f", dest="iface", type=str, default="以太网", help="扫描网卡")

# 创建ssh的解析器对象
ssh_subparsers = subparsers.add_parser("ssh", help="SSH相关漏扫")
ssh_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="SSH主机地址")
ssh_subparsers.add_argument("--port", "-P", dest="port", type=int, default=22, help="SSH端口号")
ssh_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dict/user.dict", help="SSH用户名字典文件")
ssh_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dict/passwd.dict", help="SSH密码字典文件")

# 创建Subdomain的解析器对象
subdomain_subparsers = subparsers.add_parser("subdomain", help="子域名相关漏扫")  # 修正变量名
subdomain_subparsers.add_argument("--host", "-H", dest="host", type=str, required=True, help="目标域名")
subdomain_subparsers.add_argument("--dict", "-d", dest="dict_file", type=str, default="dic.dict", help="子域名字典文件")

# 创建 port（TCP端口扫描）的解析器对象
port_subparsers = subparsers.add_parser("port", help="TCP端口扫描")
port_subparsers.add_argument("--host", "-H", dest="host", type=str, required=True, help="目标主机IP地址")
port_subparsers.add_argument("--sep", "-s", dest="sep", type=int, default=100, help="每次扫描的端口段大小 (默认 100，建议 50~200)")

# 解析命令行参数
args = parser.parse_args()

# 根据子命令调用对应功能
if args.sub_command == "mysql":
    run_mysql_for_host_or_cidr(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "redis":
    redis_run(args.host, args.port, args.password.readlines())
elif args.sub_command == "ip":
    run_ip(args.host, args.iface)
elif args.sub_command == "ssh":
    run_ssh(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "subdomain":
    run_subdomain(args.host, args.dict_file)
elif args.sub_command == "port":
    scan_port_run(args.host, args.sep)
else:
    print("请选择要使用的功能")